"""根据 java 文件所在目录加入 package 语句"""

import sys
import os
import os.path as op


def prepend(file, content, encoding='utf8'):
    with open(file, 'r+', encoding=encoding) as f:
        s = f.read()
        f.seek(0)
        f.write(content)
        f.write('\n')
        f.write(s)


def batch_package(root):
    for path, folders, files in os.walk(root):
        for file in files:
            if file.endswith('.java'):
                java_package = '.'.join(path[len(root) + 1:].split(op.sep))
                java_filename = op.realpath(op.join(path, file))
                prepend(java_filename, 'package {};'.format(java_package))


if __name__ == '__main__':
    for arg in sys.argv[1:]:
        batch_package(arg)
